{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "\n"
   ],
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "\n",
    "from torch.utils.data import Subset\n",
    "\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# 定义预处理（转换为 Tensor 并标准化）\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.1307,), (0.3081,))\n",
    "])\n",
    "\n",
    "# 下载数据集\n",
    "train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "test_data = datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n",
    "\n",
    "print(len(train_data))\n",
    "print(len(test_data))\n",
    "# train_loader = DataLoader(train_data, batch_size=64, shuffle=True)\n",
    "train_subset = Subset(train_data, indices=range(20000))\n",
    "# 创建测试集的子集（前5000个样本）\n",
    "test_subset = Subset(test_data, indices=range(8000))"
   ],
   "id": "6c12c4ca729cb914",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# 数据加载器\n",
    "train_loader = DataLoader(train_subset, batch_size=64, shuffle=True)\n",
    "test_loader = DataLoader(test_subset, batch_size=1000, shuffle=False)"
   ],
   "id": "aba142eb2e9284b0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "# print(next(enumerate(train_loader)))\n",
    "# i, (inputs, targets) = next(enumerate(train_loader))"
   ],
   "id": "fd5ccebef5d57ee0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "\n",
    "i, (inputs, targets) = next(enumerate(train_loader))\n",
    "\n",
    "# print(i)\n",
    "# print(inputs)\n",
    "# print(targets)\n",
    "# 绘图\n",
    "for i in range(25):\n",
    "    # 定义子图\n",
    "    plt.subplot(5, 5, i+1)\n",
    "    # 绘制原始像素数据\n",
    "    plt.imshow(inputs[i][0], cmap='gray')\n",
    "    # 展示图片\n",
    "plt.show()"
   ],
   "id": "cfaaaf4864a7fb26",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": "print(inputs.shape)",
   "id": "e64af98a5d2a3fe6",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Model, self).__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(28*28, 100),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(100, 10)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)\n",
    "\n",
    "model = Model().to(device)"
   ],
   "id": "de54a6c97c002535",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "epochs = 100\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "plt.ion()  # 打开交互模式\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "for i in range(epochs):\n",
    "    total_loss = 0\n",
    "    for inputs, targets in train_loader:\n",
    "        inputs, targets = inputs.to(device), targets.to(device)\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        outputs = model(inputs)\n",
    "        loss = criterion(outputs, targets)\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    if i % 1 == 0:\n",
    "        with torch.no_grad():\n",
    "            acc_accuracy = 0\n",
    "            total_samples = 0\n",
    "            for inputs, targets in test_loader:\n",
    "                inputs, targets = inputs.to(device), targets.to(device)\n",
    "                outputs = model(inputs)\n",
    "\n",
    "                outputs = outputs.argmax(dim=1)\n",
    "\n",
    "                # 计算正确率\n",
    "                acc_accuracy += (outputs == targets).sum().item()\n",
    "                total_samples += len(inputs)\n",
    "\n",
    "            avg_accuracy = acc_accuracy / total_samples\n",
    "\n",
    "            print(f'total loss: {total_loss:.3f}, accuracy:{avg_accuracy:.3f}')\n",
    "\n",
    "            ax.clear()\n",
    "            ax.plot(total_loss, label='total_loss')\n",
    "            ax.plot(avg_accuracy, label='avg_accuracy')\n",
    "            ax.set_xlabel('Epoch')\n",
    "            ax.set_ylabel('Value')\n",
    "            ax.set_title('Training Progress')\n",
    "            ax.legend()\n",
    "            plt.pause(0.1)  # 暂停0.1秒以刷新图像\n",
    "\n",
    "plt.ioff()  # 关闭交互模式\n",
    "plt.show()  # 显示最终结果"
   ],
   "id": "28fadd5fe3b15846",
   "outputs": [],
   "execution_count": null
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
